package tips.p_1000.p651_700;

import java.util.TreeMap;

/**
 * 实现一个 MapSum 类，支持两个方法，insert和 sum：
 *
 * MapSum() 初始化 MapSum 对象
 * void insert(String key, int val) 插入 key-val 键值对，
 * 字符串表示键 key ，整数表示值 val 。如果键 key 已经存在，那么原来的键值对将被替代成新的键值对。
 * int sum(string prefix) 返回所有以该前缀 prefix 开头的键 key 的值的总和。
 *
 * @author hc
 */
public class Demo677 {

    class MapSum {

        private class Node{
            public int value;
            public TreeMap<Character,Node> next;

            public Node(int value){
                this.value = value;
                next = new TreeMap<>();
            }
            public Node(){this(0);}
        }

        private Node root;

        public MapSum() {
            root = new Node();
        }

        public void insert(String key, int val) {
            Node cur = root;
            for (int i = 0; i < key.length(); i++) {
                char c = key.charAt(i);
                if (cur.next.get(c) == null){
                    cur.next.put(c,new Node());
                }
                cur = cur.next.get(c);
            }
            cur.value = val;
        }

        public int sum(String prefix) {
            Node cur = root;
            for (int i = 0; i < prefix.length(); i++) {
                char c = prefix.charAt(i);
                if (cur.next.get(c) == null){
                    return 0;
                }
                cur = cur.next.get(c);
            }
            return sum(cur);
        }

        private int sum(Node node){
            int res = node.value;
            for (char c : node.next.keySet()){
                res += sum(node.next.get(c));
            }
            return res;
        }
    }
}
